from extensions.ext_database import db
from models.robot.area_code import SystemAreaCode


def initialize_area_codes(file_path:str|None=None):
    # 存储省份、城市和区县的code，以便关联parent_code
    province_codes = {}
    city_codes = {}
    if file_path is None:
        file_path = './resource/area_code.txt'

    with open(file_path, 'r', encoding='utf-8') as file:
        for line in file:
            # 移除空白并按制表符分割
            parts = line.strip().split('\t')
            if len(parts) < 2:
                continue
            code = parts[0].strip()
            name = parts[1].strip()

            # 判断行政区划层级
            if code.endswith('0000'):  # 省级
                parent_code = None
                area_type = 1  # 假设1代表省
                name_path = name
                code_path = f"{code}000000"
                province_codes[code] = name  # 存储省份code

            elif code.endswith('00'):  # 市级
                province_code = code[:2] + '0000'  # 找到对应省的code
                parent_code = province_code
                area_type = 2  # 假设2代表市
                name_path = f"{province_codes[parent_code]}>{name}"
                code_path = f"{parent_code}>{code}000000"
                city_codes[code] = name  # 存储城市code
                parent_code = int(parent_code)*1000000

            else:  # 区县级
                city_code = code[:4] + '00'  # 找到对应市的code
                parent_code = city_code
                area_type = 3  # 假设3代表区县
                if city_code not in city_codes:
                    name_path = f"{province_codes[parent_code[:2] + '0000']}>{name}"
                    code_path = f"{parent_code[:2] + '0000'}000000>{code}000000"
                else:
                    name_path = f"{province_codes[parent_code[:2] + '0000']}>{city_codes[parent_code]}>{name}"
                    code_path = f"{parent_code[:2] + '0000'}000000>{parent_code}000000>{code}000000"
                parent_code = int(parent_code)*1000000

            # 创建 SystemAreaCode 实例
            area = SystemAreaCode(
                code=int(code)*1000000,
                name=name,
                parent_code=parent_code,
                area_type=area_type,
                name_path=name_path,
                code_path=code_path
            )
            # 添加到session
            db.session.add(area)

    # 提交事务
    db.session.commit()
    print("行政区划数据已初始化到数据库中")


class AreaService:
    @staticmethod
    def format(area_code:str)->int:
        if area_code is None:
            return 0
        if isinstance(area_code, int):
            return area_code
        if area_code.isdigit():
            if len(area_code) == 6:
                return int(f"{area_code}000000")
            elif len(area_code) == 12:
                return int(area_code)
        return -1
    
    """
    行政区划服务
    """
    @staticmethod
    def get_area_by_code(code: int) -> SystemAreaCode|None:
        """
        根据编码获取区域
        :param code: 区域编码
        :return: 区域
        """
        return db.session.query(SystemAreaCode).filter(SystemAreaCode.code == code).first()
    @staticmethod
    def get_area_name_by_code(code: int|None) -> str|None:
        """
        根据编码获取区域
        :param code: 区域编码
        :return: 区域
        """
        if code:
            area = db.session.get(SystemAreaCode,code)
            if area:
                return area.name
        return None
    
    @staticmethod
    def get_areas_by_parent(parent_code:int|None, area_type: int) -> list[SystemAreaCode]:
        """
        根据父级编码和区域类型获取区域列表
        :param parent_code: 父级编码
        :param area_type: 区域类型
        :return: 区域列表
        """
        if area_type < 1 or area_type > 3:
            return []

        if area_type == 1:
            return (
                db.session.query(SystemAreaCode)
                .filter(
                    SystemAreaCode.area_type == area_type,
                    SystemAreaCode.parent_code == 0,
                )
                .all()
            )
        else:
            return (
                db.session.query(SystemAreaCode)
                .filter(
                    SystemAreaCode.area_type == area_type,
                    SystemAreaCode.parent_code == parent_code,
                )
                .all()
            )